import torchvision

# 准备的测试数据集
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

test_data = torchvision.datasets.CIFAR100("./dataset220418", train=False, transform=torchvision.transforms.ToTensor())

# test_loader = DataLoader(dataset=test_data, batch_size=4, shuffle=True, num_workers=0, drop_last=False)  # -->step1
# test_loader = DataLoader(dataset=test_data, batch_size=64, shuffle=True, num_workers=0, drop_last=True)  # -->step2
test_loader = DataLoader(dataset=test_data, batch_size=64, shuffle=True, num_workers=0, drop_last=True)

# 测试数据集 中 第一张图片 及 target
img, target = test_data[0]
print(img.shape)
print(target)

writer = SummaryWriter("../log/dataloader")

for epoch in range(2):  # epoch会取0和1，这个是二轮洗牌！！！
    step = 0
    for data in test_loader:
        imgs, targets = data
        # print(imgs.shape)  # -->step1
        # print(targets)     # -->step1
        # writer.add_images("test_data", imgs, step)      # -->step1
        # writer.add_images("test_data_drop_last", imgs, step)        # -->step2 -->外循环for注释

        # 如果此处出现Epoch:0和Epoch:1的步数不一样，浏览器多刷新两次
        writer.add_images("Epoch: {0}".format(epoch), imgs, step)
        step = step + 1

writer.close()
